/*
CSE30332 Final Project - Iris Matcher
Group Oatmeal

Image Enhancement Program:
==========================
This program accepts an unwrapped iris image as input, enhances the image,
and writes the result to an output file. The purpose of the enhancement is
to remove average noise from the image, and to normalize the image so that
multiple iris images taken at different times will end up with the same
characteristics and intensity range. There are two different enhancement
algorithms implemented.

The first is the "averaging enhancement" method. For every pixel in the input
image, that pixel and the eight surrounding pixels are averaged together. An
intermediate image has the corresponding pixel set to that average value. The
result is an image the same size as the original, where every pixel is a blend
of the original pixel and the original surrounding pixels. Then, the
intermediate image is subtracted from the original image (each pixel in the
intermediate image has its intensity value subtracted from the intensity of
the corresponding pixel in the original image). This final image is then
scaled so that it ranges in intensity from 0-255.

The second is the "bicubic enhancement" method. This is a lightly modified
version of the Li Ma paper's enhancement method. The input image is divided
into 20x20 blocks. For each of these blocks, the average of all contained
pixels is taken. A new image is made of these averages, where one pixel in
the new image represents a 20x20 block from the original. Since our original
image is 480x40 pixels, the resultant image will be 24x2 pixels. Next, the
24x2 image is scaled back to 480x40 using a bicubic resize. This image is
then subtracted from the original image (as in the averaging method), and
the result is scaled from 0-255.

In both methods, the final image looks similar to the original, but with
local average intensities subtracted to emphasize the defining
characteristics of the iris.

We have determined experimentally that the bicubic enhancement method works
better than the averaging method. In each test case, the bicubic method was
at least marginally better than the averaging method. Thus, the program uses
the bicubic method by default.

Another feature is the ability to use a histogram version of the iris image.
This feature uses OpenCV's built-in histogram-equalization filter. Although
the output increases the image contrast and looks better-defined to a human
eye, we have found that this method worsens our matching scores in every
case we have tried. Thus, we do not use this feature by default.

In conclusion: Default options give best performance.
*/

#include <getopt.h>
#include <stdio.h>
#include <cv.h>
#include <highgui.h>
#include <string.h>

//These constants define which enhancement algorithms are available
#define ENHANCE_BICUBIC 0
#define ENHANCE_AVERAGE 1

//If debug is true, show debug messages. Gets set to true when -d given on command line
int debug = 0;

//Print usage information on the command line
void usage(char *progName)
{
	printf("Usage: %s -f INFILE [-o OUTFILE] [abgdh]\n", progName);
	printf("Input file must be a 480x40-pixel unwrapped iris\n"
		"By default, bicubic enhancement is used (-b)\n"
		"If -a is specified, use averaging enhancement method instead of bicubic\n"
		"If -g is specified, output histogram instead of the normal enhanced image\n"
		"If -o is specified, use given filename for output\n"
		"(by default, output is shown onscreen but not written to a file)\n"
		"If -d is specified, enable debug output\n"
		"-h displays this message\n");
}

//Display the given image in a window with the given title, then wait for a keypress
void showImage(char *msg, IplImage *img)
{
	cvNamedWindow(msg, 1);
	cvShowImage(msg, img);
	cvWaitKey(0);
}

//Run the Averaging Enhancement algorithm
//Give the iris image to enhance as a parameter
//Returns the enhanced image
IplImage * enhanceAverage(IplImage *src)
{
	//Declare image variables, and create images compatible with source image
	IplImage *normalized = cvCreateImage(cvGetSize(src), src->depth, 1);
	IplImage *enhanced = cvCreateImage(cvGetSize(src), src->depth, 1);
	IplImage *final = cvCreateImage(cvGetSize(src), src->depth, 1);
	//Other data variables
	int w = src->width;
	int h = src->height;
	int *valArray = malloc(w*h*sizeof(int));
	int min = 255.0, min2 = 255.0;
	int max = 0.0, max2 = 0.0;
	//Temp variables (row and column indices, etc.)
	int r, c, y, x;
	double val, avg, surroundingSum;

	//Ensure src is in the right format
	if(src->nChannels != 1)
	{
		fprintf(stderr, "src must be a 1-channel image\n");
		exit(-1);
	}

	if(debug)
		showImage("Original Image", src);
	
	//Loop through every pixel in the image
	for(c = 0; c < w; c++)
	{       
		for(r = 0; r < h; r++)
		{
			val = cvGetReal2D(src, r, c);
			
			//Keep track of the min/max pixel intensity in the source image
			if(val < min)
				min = val;
			else if(val > max)
				max = val;
			
			surroundingSum = 0.0;	//The sum of intensities around pixel (r, c)
			for(x = -2; x < 2; x++)
			{
				for(y = -2; y < 2; y++)
				{       
					surroundingSum += cvGetReal2D(src, (r+y+h)%h,(c+x+w)%w);
				}
			}
			
			//Compute average intensity of surrounding pixels
			avg = surroundingSum / 25.0;

			//The final value for this point is the original value
			//plus the average intensity of surrounding pixels
			val = val + avg;
			
			//Insert the found value into the array
			valArray[r*w + c] = (int)val;
			
			//Track the min/max pixel intensities of valArray
			if(val < min2)
				min2 = (int)val;
			if(val > max2)
				max2 = (int)val;
		}
	}
	
	//Scale the values in valArray so they range from 0-255;
	//save results into image 'enhanced'
	for(c = 0; c < w; c++)
	{
		for(r = 0; r < h; r++)
		{
			cvSetReal2D(enhanced, r, c, (double)(valArray[r*w + c] - min2) / (double)(max2 - min2) * 255.0);
		}
	}

	if(debug)
	{
		showImage("Enhanced (And Normalized)", enhanced);
		
		//If in debug mode, normalize original image and display it
		//for comparison with the enhanced version
		
		//scale all values from 0-255
		for (c = 0; c < w; c++)
		{       
			for (r = 0; r < h; r++)
			{
			      val = (cvGetReal2D(src, r, c) - min)/(max - min) * 255;
			      cvSetReal2D(normalized, r, c, val);
			}
		}
		showImage("Normalized Original", normalized);
	}
		        
	return enhanced;
}

//Run the Bicubic Enhancement algorithm
//Give the iris image to enhance as a parameter
//Returns the enhanced image
IplImage * enhanceBicubic(IplImage* src)
{
	//Allocate space for temporary images
	//'temp' will be 1/20th the size of the source image;
	//each pixel in temp will be the average intensity of a 20x20 block from src
	IplImage* temp = cvCreateImage(cvSize(src->width/20,src->height/20), src->depth, 1);
	//'res' will be 'temp' resized back to the orignal source image dimensions
	IplImage* res = cvCreateImage(cvGetSize(src), src->depth, 1);
	//'sub' will contain the result of subtracting (src - temp)
	IplImage* sub = cvCreateImage(cvGetSize(src), src->depth, 1);
	int tval;
	double avg, val, min, max;
	//Row and column for the temporary, block, and source coordinates
	int cTemp = 0, cBlock = 0, c = 0;
	int rTemp = 0, rBlock = 0, r = 0;

	if(debug)
		showImage("Original Image", src);
	
	//Loop over 20x20 blocks in the source image, average their
	//intensities, and make a new image 'temp' where each pixel
	//is one of these averages
	for(c = 0; c < src->width; c += 20)
	{
	        for(r = 0; r < src->height; r += 20)
	        {
			tval = 0;
			for(cBlock = 0; cBlock < 20; cBlock++)
			{
				for(rBlock = 0; rBlock < 20; rBlock++)
				{
					tval += (int)cvGetReal2D(src, r+rBlock, c+cBlock);
				}
			}

			avg = tval / 400;
			if(debug)
				printf("Avg for (%d,%d):%d->(%d,%d)\n", c, r, avg, cTemp, rTemp);
			cvSetReal2D(temp, rTemp, cTemp, (double)avg);
			rTemp++;
		}

		cTemp++;
		rTemp=0;
	}
	
	if(debug)
		showImage("Avg Squares", temp);

	//Scale up the temp image to be the size of the source image
	cvResize(temp, res, CV_INTER_CUBIC);
	
	if(debug)
		showImage("Resized", res);

	//Find min and max values
	min = 255;	//Set original min to the max. possible value
	max = 0;	//Set original max to the min. possible value
	for(c = 0; c < src->width; c++)
	{
		for (r = 0; r < src->height; r++)
		{
			val = cvGetReal2D(src, r, c) - cvGetReal2D(res, r, c);

			if(val < min)
				min = val;
			else if(val > max)
				max = val;
		}
	}

	//Fix min and max in the special case that the image was one solid color
	if(min > max)
		if(min == 0.0)
			min = max;
		else
			max = min;
	
	//Perform subtraction of (src-res) and store results in 'sub'
	for(c = 0; c < src->width; c++)
	{
		for(r = 0; r < src->height; r++)
		{
			val = cvGetReal2D(src, r, c) - cvGetReal2D(res, r, c);
			val = (val-min) / (max-min) * 255.0;
			cvSetReal2D(sub, r, c, val);
		}
	}

	if(debug)
		showImage("After Subtraction", sub);
	
	return sub;
}

int main(int argc, char * argv[])
{
	char *fname = NULL;
	char *ofile = NULL;
	char c;
	int enhanceMethod = ENHANCE_BICUBIC;
	int histo = 0;
	IplImage *src, *outImage;
	IplImage *hist;


	//handle the various arguments (only hostname)
	while((c = getopt(argc, argv, "?df:abho:")) != -1)
	{
		switch(c)
		{
			case 'f':
				fname = optarg;
				break;
			case 'a':
				enhanceMethod = ENHANCE_AVERAGE;
				break;
			case 'b':
				enhanceMethod = ENHANCE_BICUBIC;
				break;
			case 'o':
				ofile = optarg;
				break;
			case 'g':
				histo++;
				break;
			case 'd':
				debug = 1;
				break;
			case 'h':
				usage(*argv);
				exit(0);
			default:
				fprintf(stderr, "Error; Need filename, method\n");
		}
	}

	if (fname == NULL)
	{
		fprintf(stderr, "Must specify a valid input filename\n");
		usage(*argv);
		exit(-1);
	}

	//Read in the source image
	src = cvLoadImage(fname, 0);
	if(!src)
	{
		fprintf(stderr, "OpenCV could not load image file '%s'; aborting\n", fname);
		exit(-1);
	}

	if(src->width != 480 || src->height != 40)
	{
		fprintf(stderr, "Input file must be a 480x40-pixel unwrapped iris image");
		exit(-1);
	}

	//Run chosen enhancement method
	if (enhanceMethod == ENHANCE_AVERAGE)
		outImage = enhanceAverage(src);
	else if (enhanceMethod == ENHANCE_BICUBIC)
		outImage = enhanceBicubic(src);
	
	//Run the histogram equalizer
	if(histo)
	{
		hist = cvCreateImage(cvGetSize(src), src->depth, 1);
		cvEqualizeHist(outImage, hist);
	}

	if(debug && histo)
		showImage("Hist", hist);

	//If requested, save the result to an output file
	//Otherwise, just show the image onscreen
	if(ofile)
	{
		if(histo)
			cvSaveImage(ofile, hist);
		else
			cvSaveImage(ofile, outImage);
	}
	else if(!debug)	//If in debug mode, the output was already displayed
	{
		if(histo)
			showImage("Output", hist);
		else
			showImage("Output", outImage);
	}

	return 0;
}

